# -*- coding: utf-8 -*-
# @Author  : Assistant
# @Time    : 2025/1/27
# @Function: 批量更新QA_final目录中Excel文件的category_id字段

import os
import pandas as pd
from pathlib import Path

# 水库名称与category_id的映射关系
RESERVOIR_MAPPING = {
    "白石水库调度规程": "21200510",
    "清河水库调度规程": "20810200",
    "石佛寺水库调度规程": "20600340",
    "大伙房水库调度规程": "21100150",
    "观音阁水库调度规程": "21110150",
    "葠窝水库调度规程": "21110400",
    "柴河水库调度规程": "20800900",
    "汤河水库调度规程": "21113800",
    "闹德海水库调度规程": "20910930"
}

def update_category_id_in_excel(file_path: str, category_id: str):
    """更新Excel文件中的category_id字段
    
    Args:
        file_path (str): Excel文件路径
        category_id (str): 新的category_id值
    """
    try:
        # 读取Excel文件
        df = pd.read_excel(file_path)
        
        # 检查是否存在category_id列
        if 'category_id' not in df.columns:
            # 如果不存在，添加category_id列
            df['category_id'] = category_id
            print(f"  添加category_id列: {category_id}")
        else:
            # 如果存在，更新所有行的category_id
            df['category_id'] = category_id
            print(f"  更新category_id: {category_id}")
        
        # 保存更新后的文件
        df.to_excel(file_path, index=False)
        print(f"  文件已保存: {os.path.basename(file_path)}")
        
        return len(df)
        
    except Exception as e:
        print(f"  错误：处理文件 {os.path.basename(file_path)} 时出错: {e}")
        return 0

def extract_reservoir_name(filename: str) -> str:
    """从文件名提取水库名称
    
    Args:
        filename (str): 文件名（包含.xlsx扩展名）
        
    Returns:
        str: 水库名称（不含扩展名）
    """
    # 移除.xlsx扩展名
    return os.path.splitext(filename)[0]

def main():
    """主函数 - 批量更新category_id"""
    print("=== QA文件category_id批量更新工具 ===")
    print()
    
    # 设置QA_final目录路径
    current_dir = os.path.dirname(os.path.abspath(__file__))
    qa_final_dir = os.path.join(current_dir, "QA_final")
    
    print(f"当前工作目录: {current_dir}")
    print(f"QA_final目录: {qa_final_dir}")
    print()
    
    # 检查目录是否存在
    if not os.path.exists(qa_final_dir):
        print(f"错误：QA_final目录不存在 - {qa_final_dir}")
        return
    
    # 获取所有Excel文件
    excel_files = [f for f in os.listdir(qa_final_dir) if f.endswith('.xlsx')]
    
    if not excel_files:
        print(f"在 {qa_final_dir} 中未找到任何Excel文件")
        return
    
    print(f"找到 {len(excel_files)} 个Excel文件")
    print("\n水库ID映射关系:")
    for reservoir, category_id in RESERVOIR_MAPPING.items():
        print(f"  {reservoir}: {category_id}")
    print()
    
    # 统计信息
    total_files = 0
    updated_files = 0
    total_records = 0
    skipped_files = []
    
    # 处理每个Excel文件
    for excel_file in excel_files:
        total_files += 1
        file_path = os.path.join(qa_final_dir, excel_file)
        
        # 提取水库名称
        reservoir_name = extract_reservoir_name(excel_file)
        
        print(f"正在处理第 {total_files}/{len(excel_files)} 个文件: {excel_file}")
        print(f"  水库名称: {reservoir_name}")
        
        # 查找对应的category_id
        if reservoir_name in RESERVOIR_MAPPING:
            category_id = RESERVOIR_MAPPING[reservoir_name]
            print(f"  匹配到category_id: {category_id}")
            
            # 更新文件
            record_count = update_category_id_in_excel(file_path, category_id)
            if record_count > 0:
                updated_files += 1
                total_records += record_count
                print(f"  成功更新 {record_count} 条记录")
        else:
            print(f"  警告：未找到匹配的category_id，跳过文件")
            skipped_files.append(excel_file)
        
        print()
    
    # 输出处理结果统计
    print("=== 处理完成 ===")
    print(f"总文件数: {total_files}")
    print(f"成功更新: {updated_files}")
    print(f"跳过文件: {len(skipped_files)}")
    print(f"总记录数: {total_records}")
    
    if skipped_files:
        print("\n跳过的文件:")
        for skipped_file in skipped_files:
            print(f"  - {skipped_file}")
    
    print("\n所有文件的category_id已按照映射关系更新完成！")

if __name__ == "__main__":
    main()